# !/usr/bin/python3
# -*- coding:utf-8 -*-
# Copyright 2021 The Chinaunicom Software Team. All rights reserved.
# @FileName: annoy.py
# @Author  : Dammy
# @Time    : 2021/12/17



from annoy import AnnoyIndex
import random
import numpy as np
import json, re, os, codecs, time
from tqdm import tqdm
from src.intelligent_interaction.engine.sem_conf import Configuration
# from logserver import LogServer
config = Configuration()

### Log ###
# log_server = LogServer(app='sqr', log_path="./log_sqr/annoy_log/")

class AnnoyRecall:
    def __init__(self, f_dim, all_tid=False):
        """
        :paramter f_dim  : 向量维度初始化
        :paramter n_sims :
        :paramter all_tid: 是否输入所有的tid, 默认为False
        :paramter npy_data: 是否输入npy数据, 默认为False
        """
        self.f_dim = f_dim
        self.sim = AnnoyIndex(self.f_dim, 'angular')
        ann_path = config.LoadParams['ann']

#         if all_tid:
#             self.all_tid = all_tid
#         else:
#             self.all_tid = list(self.origin_all_data_json.keys())

        if os.path.exists(ann_path):
            self.sim.load(ann_path, prefault=False)


    def load_npy(self):
        """
        load .npy data
        :paramter path : 待载入路径
        """
        npy_data = np.load(self.npy_path)
        print("加载npy完成..., 共 {} 条".format(len(npy_data)))
        print("=========================================" + "\n")
        return npy_data


    def train_annoy_index_online(self, all_id, npy_data, ann_path):
        """
        inner method: 实时在线训练annoy索引
        :paramter :
        """
        assert len(all_id) == len(npy_data)  # The length of the generated npy file is not the same as the tid length

        t = AnnoyIndex(self.f_dim, 'angular')   # 实例化annoy, 用余弦相似度
        for i in range(len(all_id)):
            # print("第 {} 个ing...".format(i))
            # v = [random.gauss(0, 1) for z in range(f_dim)]
            t.add_item(i, npy_data[i])
        t.build(150)   # 建立n_trees树的森林，查询时，树越多，精度越高.
        t.save(ann_path, prefault=False)
        print("Save ANN Finish!")


    def index_recall(self, ori_tid, ori_type, ori_labels, index, n_sims, search_k, visualization=False):
        """
        Inner Method: 使用tid的index召回相似题方法（可以选的多个或单个）
        :paramter index :  tid映射的index
        :paramter n_sims:  召回的数据
        :paramter search_k:  搜索权衡因子
        """
        start_time = time.time()

        # 批量测试
        # lis = [self.sim.get_nns_by_item(i+1, 100) for i in range(100)]
        # for i in range(10000):
        #     u.get_nns_by_item(i+1, 100)
        # _index = self.all_tid.index(index)

        sim_subjec_info = []  # 最后返回的结果
        result = self.sim.get_nns_by_item(i=index, n=n_sims, search_k=search_k, include_distances=True)

        start_time_2 = time.time()
        for i in range(len(result[0])):
            tid = self.all_tid[result[0][i]]
            weight = result[1][i]

            if tid != ori_tid:
                if tid not in self._dup:
                    sub_type = self.origin_all_math_data_json[tid]['type']
                    elite    = self.origin_all_math_data_json[tid]['elite']
                    info = (tid, self.origin_all_math_data_json[tid]['labels'], elite, sub_type, self.origin_all_math_data_json[tid]['question'])
                    # info = "第{}道, tid: {}, 知识点: {}, 精品题: {}, 题目类型: {}, 相似题: {}".format(num, tid, self.origin_all_math_data_json[tid]['labels'], elite, sub_type, self.origin_all_math_data_json[tid]['question'])
                    sim_subjec_info.append(info)
                    # print("第{}道: 相似题:{}, labels: {}, weight: {}".format(i+1, self.all_math_data_json[tid]['question'], ",".join(self.all_math_data_json[tid]['labels']), weight))
        end_time_2 = time.time()
        print("Annoy Time 2:", end_time_2 - start_time_2)
        return sim_subjec_info


    def vecter_recall(self, vector, n_sims, search_k):
        """
        Inner Method: 没有重排序的简单的vector召回
        :paramter ori_tid      :
        """
        sims_info = []  # 最后返回的结果
        result = self.sim.get_nns_by_vector(vector, n_sims, search_k, include_distances=True)
        qids, distance = result[0], result[1]
        _distance = [1-i for i in distance]
        # print("ANNOY result:", result, _distance)
        return qids, _distance 

